#include "XgModulesCollector.h"
#include <tchar.h>
#include <sys\types.h> 
#include <sys\stat.h> 

typedef struct {  
	ULONG i[2];  
	ULONG buf[4];  
	unsigned char in[64];  
	unsigned char digest[16];
} MD5_CTX;

HMODULE g_hCryptDll;

typedef void (WINAPI *MD5Init)(MD5_CTX* /*context */);
typedef void (WINAPI *MD5Update)(MD5_CTX* /*context*/,unsigned char* /*input*/,unsigned int /*inlen*/);
typedef void (WINAPI *MD5Final)(MD5_CTX* /*context*/);

MD5Init g_fMD5Init;
MD5Update g_fMD5Update;
MD5Final g_fMD5Final;

XgModulesCollector::XgModulesCollector(XgW32Processes processes)
{
	g_hCryptDll = LoadLibrary(_T("cryptdll.dll"));
	g_fMD5Init = (MD5Init)GetProcAddress(g_hCryptDll, "MD5Init");
	g_fMD5Update = (MD5Update)GetProcAddress(g_hCryptDll, "MD5Update");
	g_fMD5Final = (MD5Final)GetProcAddress(g_hCryptDll, "MD5Final");
	
	collect(processes);
}

XgModulesCollector::~XgModulesCollector(void)
{
	FreeLibrary(g_hCryptDll);
}

bool XgModulesCollector::collect(XgW32Processes processes)
{
	for( XgW32Processes::Iterator iter = processes.begin(); iter != processes.end(); iter++ )
	{
		DWORD pid = iter->first;
		XgW32Process process = iter->second;
		
		XgW32ProcessModules modules = process.getModules();
		for( XgW32ProcessModules::Iterator iter2 = modules.begin(); iter2 != modules.end(); iter2++ )
		{
			XgW32ProcessModule proc_module = iter2->second;
			
			//int s;

			////const TCHAR* szPath = proc_module.path().c_str();

			//struct _stat fileStat;
			//int err = _tstat( proc_module.path().c_str(), &fileStat );
			//if( err != 0 ) 
			//	s = 0;
			//
			//s = fileStat.st_size;

			//HANDLE hFile = CreateFile(proc_module.path().c_str(), 
			//	FILE_READ_DATA, 
			//	FILE_SHARE_READ, 
			//	NULL,
			//	OPEN_EXISTING,
			//	FILE_ATTRIBUTE_NORMAL,
			//	NULL
			//	);
			//
			//char* lpBuffer = new char[s];

			//DWORD cBytes;

			//if( hFile != NULL )
			//{
			//	ReadFile(hFile, lpBuffer, s, &cBytes, NULL);
			//}
			//
			//MD5_CTX md5ctx;

			//g_fMD5Init(&md5ctx);
			//g_fMD5Update(&md5ctx, (unsigned char*)lpBuffer, s);
			//g_fMD5Final(&md5ctx);
			//
			//TCHAR digest[16];
			//
			//#ifdef _UNICODE
			//xg_char2wchar(digest, 16, (char*)md5ctx.digest, 16);
			//#else
			//strcpy_s(digest, 16, (char*)md5ctx.digest);
			//#endif
			//
			//
			//TCHAR szDigest[32];

			//for( int i = 0; i < 16; i++ )
			//{
			//	_itot_s( ((BYTE*)digest)[i], &szDigest[i], 3, 16 );
			//}

			////TSTRING strDigest(digest);
			////_map[strDigest] = XgW32Module(proc_module.name());

			//delete [] lpBuffer;
			//CloseHandle(hFile);

			char szPath[MAX_PATH];

			size_t len = _tcslen(proc_module.path().c_str());

#ifdef _UNICODE
			xg_wchar2char( szPath, MAX_PATH, proc_module.path().c_str(), len + 1 );
#else
			strcpy_s( szPath, MAX_PATH, proc_module.path().c_str() );
#endif

			MD5_CTX md5ctx;

			g_fMD5Init(&md5ctx);
			g_fMD5Update(&md5ctx, (unsigned char*)szPath, (int)len);
			g_fMD5Final(&md5ctx);

			TCHAR szDigest[33];
			szDigest[32] = '\0';
			TCHAR szTemp[8];
			int curr_pos = 0;
	
			for( int i = 0; i < 16; i++ )
			{
				_itot_s(md5ctx.digest[i], szTemp, 8, 16);
				len = _tcslen(szTemp);
				if( len == 1 )
				{
					szTemp[1] = szTemp[0];
					szTemp[0] = '0';
					szTemp[2] = '\0';
					len = 2;
				}

				for( int j = 0; j < (int)len; j++ )
				{
					szDigest[curr_pos + j] = szTemp[j];
				}
				curr_pos = curr_pos + (int)len;
			}

			TSTRING strDigest(szDigest);
			_map[strDigest] = XgW32Module( proc_module.name(), proc_module.path() );
		}
	}
	return true;
}